import pandas as pd
import re
from gq.trec_qa.trec_const import Trec_Const

baidu_stopwords = [line.strip() for line in open(Trec_Const.baidu_stop_words, encoding="utf-8").readlines()]  # 加载停用词


def load_data():
    train = pd.read_csv(Trec_Const.train_path, sep=":", header=None, encoding="gbk",
                        names=["label", "question"])
    test = pd.read_csv(Trec_Const.test_path, sep=":", header=None, encoding="gbk",
                       names=["label", "question"])
    return train, test


def clean_text(text):
    pat_letter = re.compile(r'[^a-zA-Z \ ]+')
    text = pat_letter.sub('', text)
    return text


def process():
    train, test = load_data()
    train['question'] = train["question"].astype('str')
    train["question"] = train["question"].apply(clean_text)
    train['question'] = train['question'].str.lower().str.split()
    train['question'] = train['question'].map(lambda x: x[1:])
    train['question'] = train['question'].apply(lambda x: [item for item in x if item not in baidu_stopwords])
    test['question'] = test["question"].astype('str')
    test["question"] = test["question"].apply(clean_text)
    test['question'] = test['question'].str.lower().str.split()
    test['question'] = test['question'].map(lambda x: x[1:])
    test['question'] = test['question'].apply(lambda x: [item for item in x if item not in baidu_stopwords])
    train.to_csv('./data/train.csv')
    test.to_csv('./data/test.csv')


def load_saved_data():
    train_data = pd.read_csv(Trec_Const.train_saved_path, encoding='utf-8')
    test_data = pd.read_csv(Trec_Const.test_saved_path, encoding='utf-8')
    x_train = train_data["question"]
    y_train = train_data["label"]
    x_test = test_data["question"]
    y_test = test_data["label"]

    return x_train, y_train, x_test, y_test


if __name__ == '__main__':
    print(load_saved_data())
